import torch
import torch.nn as nn
import torch.nn.functional as F
from model.modules.atst.atst_model import ATST
from model.modules.classifier import BaseCls
from model.modules.cnn import CNN
from model.modules.rnn import RNN

class CRNN(nn.Module):
    def __init__(
        self,
        cnn: CNN,
        rnn: RNN,
        classifier: BaseCls,
        atst_dropout=0.0,
        atst_train=False
    ):
        super(CRNN, self).__init__()
        self.cnn = cnn
        self.rnn = rnn
        self.classifier = classifier
        self.atst_dropout = atst_dropout
        self.init_atst()
        self.merge_layer = torch.nn.Linear(cnn.channels[-1] + self.atst_frame.atst.embed_dim, cnn.channels[-1])        
        if atst_train:
            for param in self.atst_frame.parameters():
                param.requires_grad = True

    def init_atst(self, path=None):
        if path is None:
            atst_path = "./pretrained_ckpts/atstframe_as2M.ckpt"
        else:
            atst_path = path
        print("Loading ATST from:", atst_path)
        self.atst_frame = ATST(atst_path, atst_dropout=self.atst_dropout)
        self.atst_frame.eval()
        for param in self.atst_frame.parameters():
            param.detach_()

    def forward(self, x):
        x, pretrain_x = x[0], x[1]
        # CNN
        x = x.transpose(1, 2).unsqueeze(1)
        x = self.cnn(x)
        bs, chan, frames, freq = x.size()
        x = x.permute(0, 2, 1, 3).reshape(bs, frames, -1)
        # ATST-Frame
        pretrain_x = self.atst_frame(pretrain_x)
        # Merge  
        pretrain_x = F.adaptive_avg_pool1d(pretrain_x, x.shape[-2]).transpose(1, 2)
        x = self.merge_layer(torch.cat((x, pretrain_x), -1))
        # RNN
        x = self.rnn(x)
        strong, weak = self.classifier(x)
        return strong, weak
